Ground state simulation of the quantum Ising model.

# pylint: disable=invalid-name

# pylint: disable-next=import-error
from mpi4py import MPI

import qtealeaves as qtl
from qtealeaves.emulator import TTN
from qtealeaves.models import get_quantum_ising_1d
from qtealeaves.tensors import TensorBackend

# pylint: disable-next=c-extension-no-member
comm = MPI.COMM_WORLD


def groundstate(tn_type=5):
    """
    Main method for the ground state simulation of 1d quantum
    Ising model.

    **Arguments**

    tn_type : int, optional
        Choose 5 for python-TTN, 6 for python-MPS.
        Default to 5.

    """
    input_folder = lambda params: "QI1d/input_L%d" % (params["L"])
    output_folder = lambda params: "QI1d/output_L%d" % (params["L"])

    model, my_ops = get_quantum_ising_1d()

    my_conv = qtl.convergence_parameters.TNConvergenceParameters(
        max_iter=7, max_bond_dimension=20
    )

    my_obs = qtl.observables.TNObservables()
    my_obs += qtl.observables.TNState2File("QI1d/psi", "U")

    simulation = qtl.QuantumGreenTeaSimulation(
        model,
        my_ops,
        my_conv,
        my_obs,
        tn_type=tn_type,
        folder_name_input=input_folder,
        folder_name_output=output_folder,
        has_log_file=False,
        store_checkpoints=False,
    )

    params = []

    params.append({"L": 8, "J": 1.0, "g": 0.5})
    simulation.run(params, delete_existing_folder=True)

    return TTN.read("QI1d/psi.pklttn", TensorBackend())


def main(last_bit="0"):
    """Calculate groundstate on MPI-thread 0 and then sample in parallel."""
    if comm.Get_rank() == 0:
        psi = groundstate()
    else:
        psi = None

    comm.Barrier()

    py_tensor_backend = TensorBackend()

    # pylint: disable-next=unused-argument
    def filter_func(state, interval, last_bit=last_bit):
        return state[-1] == last_bit

    samples = TTN.mpi_sample_n_unique_states(
        psi,
        100,
        comm,
        py_tensor_backend,
        filter_func=filter_func,
        mpi_final_op="mpi_all_gather",
    )

    print("Rank & samples", comm.Get_rank(), len(samples))


if __name__ == "__main__":
    main()